import copy
import heapq
from typing import List


class Solution:
    def trapRainWater(self, heightMap: List[List[int]]) -> int:
        m = len(heightMap)
        n = len(heightMap[0])
        heap = []
        visited = [[False]*n for _ in range(m)]

        # 边界入堆
        for i in range(m):
            heapq.heappush(heap, (heightMap[i][0], i, 0))
            heapq.heappush(heap, (heightMap[i][n-1], i, n-1))
        for j in range(1, n-1):
            heapq.heappush(heap, (heightMap[0][j], 0, j))
            heapq.heappush(heap, (heightMap[m-1][j], m-1, j))

        ans = 0
        directions = [(1,0),(-1,0),(0,1),(0,-1)]

        while heap:
            h, x, y = heapq.heappop(heap)
            if visited[x][y]: continue
            visited[x][y] = True
            ans += max(0, h - heightMap[x][y])
            for dx, dy in directions:
                nx, ny = x+dx, y+dy
                if 0 <= nx < m and 0 <= ny < n and not visited[nx][ny]:
                    heapq.heappush(heap, (max(heightMap[nx][ny], h), nx, ny))
        return ans
if __name__ == '__main__':
    s = Solution()
    r = s.trapRainWater([[1,4,3,1,3,2],[3,2,1,3,2,4],[2,3,3,2,3,1]])
    print(r)